import argparse
import threading
import pymysql
import redis
import ipaddress
import warnings
import socket
import paramiko
from scapy.layers.l2 import ARP
from scapy.sendrecv import sr1
from urllib.parse import urlparse
from concurrent.futures import ThreadPoolExecutor


# 颜色定义类
class Colors:
    RESET = "\033[0m"
    BLACK = "\033[30m"
    RED = "\033[31m"
    GREEN = "\033[32m"
    YELLOW = "\033[33m"
    BLUE = "\033[34m"
    MAGENTA = "\033[35m"
    CYAN = "\033[36m"
    WHITE = "\033[37m"
    BLACK_BRIGHT = "\033[90m"
    RED_BRIGHT = "\033[91m"
    GREEN_BRIGHT = "\033[92m"
    YELLOW_BRIGHT = "\033[93m"
    BLUE_BRIGHT = "\033[94m"
    MAGENTA_BRIGHT = "\033[95m"
    CYAN_BRIGHT = "\033[96m"
    WHITE_BRIGHT = "\033[97m"
    BG_BLACK = "\033[40m"
    BG_RED = "\033[41m"
    BG_GREEN = "\033[42m"
    BG_YELLOW = "\033[43m"
    BG_BLUE = "\033[44m"
    BG_MAGENTA = "\033[45m"
    BG_CYAN = "\033[46m"
    BG_WHITE = "\033[47m"
    BOLD = "\033[1m"
    UNDERLINE = "\033[4m"
    REVERSE = "\033[7m"


# 程序横幅
banner = """
____              ___ ____    ___________  ____     ___ ____     ______     
`Mb(      db      )d' `MM'    `MM`MMMMMMMb.`Mb(     )d' `MM'     `M`MM'     
 YM.     ,PM.     ,P   MM      MM MM    `Mb YM.     ,P   MM       M MM      
 `Mb     d'Mb     d'   MM      MM MM     MM `Mb     d'   MM       M MM      
  YM.   ,P YM.   ,P    MM      MM MM     MM  YM.   ,P    MM       M MM      
  `Mb   d' `Mb   d'    MMMMMMMMMM MM    .M9  `Mb   d'    MM       M MM      
   YM. ,P   YM. ,P     MM      MM MMMMMMM9'   YM. ,P     MM       M MM      
   `Mb d'   `Mb d'     MM      MM MM  \M\     `Mb d'     MM       M MM      
    YM,P     YM,P      MM      MM MM   \M\     YM,P      YM       M MM      
    `MM'     `MM'      MM      MM MM    \M\    `MM'       8b     d8 MM    / 
     YP       YP      _MM_    _MM_MM_    \M\_   YP         CZY _(v1.3) 
                  功能：MySQL+Redis+SSH弱口令 | IP存活 | 子域名扫描
"""
print(f"{Colors.BLUE_BRIGHT}{Colors.BOLD}{banner}{Colors.RESET}")


# -------------------------- 1. MySQL扫描 --------------------------
def scan_mysql(host="127.0.0.1", port=3306, user="root", password="", connect_timeout=5):
    conn = None
    try:
        conn = pymysql.connect(
            host=host,
            port=port,
            user=user.strip(),
            password=password.strip(),
            connect_timeout=connect_timeout,
            database="information_schema",
            charset="utf8mb4",
            autocommit=True,
        )
        print(f"{Colors.GREEN}[+] MySQL  {host}:{port}  弱口令 -> {user.strip()} / {password.strip()}{Colors.RESET}")
    except pymysql.OperationalError:
        pass
    except Exception as e:
        print(f"{Colors.RED}[-] MySQL  {host}:{port}  异常 -> {str(e)}{Colors.RESET}")
    finally:
        if conn:
            conn.close()


def mysql_run(host, port, users, passwds):
    try:
        with ThreadPoolExecutor(max_workers=50) as executor:
            for user in users:
                for passwd in passwds:
                    executor.submit(scan_mysql, host, port, user, passwd)
    except Exception as e:
        print(f"{Colors.RED}[-] MySQL模块异常: {str(e)}{Colors.RESET}")


# -------------------------- 2. Redis扫描 --------------------------
def scan_redis(host="127.0.0.1", port=6379, password="", timeout=5):
    redis_client = None
    try:
        redis_kwargs = {
            "host": host,
            "port": port,
            "db": 0,
            "socket_timeout": 3,
            "socket_connect_timeout": timeout,
            "decode_responses": False
        }
        if password.strip():
            redis_kwargs["password"] = password.strip()

        redis_client = redis.Redis(**redis_kwargs)
        if redis_client.ping():
            if not password.strip():
                print(f"{Colors.GREEN}[+] Redis  {host}:{port}  未授权访问{Colors.RESET}")
            else:
                print(f"{Colors.GREEN}[+] Redis  {host}:{port}  弱口令 -> {password.strip()}{Colors.RESET}")
    except redis.AuthenticationError:
        pass
    except redis.ConnectionError:
        pass
    except Exception as e:
        print(f"{Colors.RED}[-] Redis  {host}:{port}  异常 -> {str(e)}{Colors.RESET}")
    finally:
        if redis_client:
            redis_client.close()


def redis_run(host, port, passwds):
    try:
        with ThreadPoolExecutor(max_workers=30) as executor:
            for passwd in passwds:
                executor.submit(scan_redis, host, port, passwd)
    except Exception as e:
        print(f"{Colors.RED}[-] Redis模块异常: {str(e)}{Colors.RESET}")


# -------------------------- 3. IP存活扫描 --------------------------
def arp_scan_ip(pdst="127.0.0.1", iface="以太网"):
    try:
        send_pkg = ARP(pdst=pdst)
        result_pkg = sr1(send_pkg, iface=str(iface), verbose=False, timeout=0.2)
        if result_pkg:
            psrc = result_pkg.sprintf("%ARP.psrc%")
            pmack = result_pkg.sprintf("%ARP.hwsrc%")
            print(f"{Colors.GREEN}[+] 存活主机  {psrc}  |  MAC: {pmack}{Colors.RESET}")
    except Exception as e:
        print(f"{Colors.RED}[-] ARP扫描  {pdst}  异常 -> {str(e)}{Colors.RESET}")


def tcp_ping(ip="127.0.0.1", port=80, timeout=1):
    sock = None
    try:
        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        sock.settimeout(timeout)
        sock.connect((ip, port))
        print(f"{Colors.GREEN}[+] 存活主机  {ip}  |  端口{port}开放{Colors.RESET}")
    except (socket.timeout, ConnectionRefusedError):
        pass
    except Exception as e:
        print(f"{Colors.RED}[-] TCP扫描  {ip}  异常 -> {str(e)}{Colors.RESET}")
    finally:
        if sock:
            sock.close()


def run_ip(host, iface, port=80):
    warnings.filterwarnings("ignore")
    try:
        network = ipaddress.IPv4Network(host, strict=False)
        print(f"{Colors.YELLOW}[*] 扫描网段: {network} （ARP+TCP{port}检测）{Colors.RESET}")

        with ThreadPoolExecutor(max_workers=100) as executor:
            for ip in network:
                ip_str = str(ip)
                executor.submit(arp_scan_ip, ip_str, iface)
                executor.submit(tcp_ping, ip_str, port)
    except ValueError as e:
        print(f"{Colors.RED}[-] IP解析错误: {str(e)}（正确格式：192.168.1.0/24）{Colors.RESET}")
    except Exception as e:
        print(f"{Colors.RED}[-] IP模块异常: {str(e)}{Colors.RESET}")


# -------------------------- 4. SSH扫描 --------------------------
def scan_ssh(host="127.0.0.1", port=22, user="root", password="", timeout=10):
    ssh_client = None
    try:
        ssh_client = paramiko.SSHClient()
        ssh_client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        # 1. 新增 banner_timeout，避免读取SSH协议头超时
        ssh_client.connect(
            hostname=host,
            port=port,
            username=user.strip(),
            password=password.strip(),
            timeout=timeout,  # 连接超时
            banner_timeout=timeout,  # 读取SSH协议头超时（关键新增）
            allow_agent=False,
            look_for_keys=False
        )
        print(f"{Colors.GREEN}[+] SSH  {host}:{port}  弱口令 -> {user.strip()} / {password.strip()}{Colors.RESET}")

    # 2. 捕获认证失败（不提示，避免刷屏）
    except paramiko.AuthenticationException:
        pass

    # 3. 捕获连接超时/被拒绝（简洁提示）
    except socket.timeout:
        print(f"{Colors.RED}[-] SSH  {host}:{port}  连接超时（{timeout}秒）{Colors.RESET}")
    except ConnectionRefusedError:
        print(f"{Colors.RED}[-] SSH  {host}:{port}  连接被拒绝（服务未启动？）{Colors.RESET}")

    # 4. 捕获靶机主动关闭连接（10054错误，简洁提示，不打印堆栈）
    except (ConnectionResetError, paramiko.ssh_exception.SSHException) as e:
        if "10054" in str(e) or "Error reading SSH protocol banner" in str(e):
            # 只输出简短提示，避免堆栈刷屏
            print(f"{Colors.RED}[-] SSH  {host}:{port}  连接被靶机临时关闭（防暴力破解）{Colors.RESET}")
        else:
            print(f"{Colors.RED}[-] SSH  {host}:{port}  SSH协议异常 -> {str(e)}{Colors.RESET}")

    # 5. 其他未知异常
    except Exception as e:
        print(f"{Colors.RED}[-] SSH  {host}:{port}  未知异常 -> {str(e)}{Colors.RESET}")

    # 6. 确保连接关闭（避免资源泄漏）
    finally:
        if ssh_client:
            try:
                ssh_client.close()
            except:
                pass


def ssh_run(host, port, users, passwds):
    try:
        print(f"{Colors.YELLOW}[*] SSH扫描: {host}:{port}（字典加载中）{Colors.RESET}")
        with ThreadPoolExecutor(max_workers=40) as executor:
            for user in users:
                for passwd in passwds:
                    executor.submit(scan_ssh, host, port, user, passwd)
    except Exception as e:
        print(f"{Colors.RED}[-] SSH模块异常: {str(e)}{Colors.RESET}")


# -------------------------- 5. 子域名扫描 --------------------------
def resolve_domain(domain, timeout=3):
    try:
        ip_list = socket.gethostbyname_ex(domain)[2]
        for ip in ip_list:
            print(f"{Colors.GREEN}[+] 子域名  {domain}  -> IP: {ip}{Colors.RESET}")
    except socket.gaierror:
        pass
    except Exception as e:
        print(f"{Colors.RED}[-] 子域名  {domain}  异常 -> {str(e)}{Colors.RESET}")


def subdomain_run(target, subdomain_dict, timeout=3):
    try:
        parsed = urlparse(target)
        if parsed.netloc:
            target_domain = parsed.netloc.replace("www.", "")
        else:
            target_domain = target.replace("www.", "")

        print(f"{Colors.YELLOW}[*] 子域名扫描: 主域名={target_domain}{Colors.RESET}")

        subdomains = [line.strip() for line in subdomain_dict if line.strip()]
        if not subdomains:
            print(f"{Colors.RED}[-] 子域名字典为空！{Colors.RESET}")
            return

        with ThreadPoolExecutor(max_workers=80) as executor:
            for sub in subdomains:
                full_domain = f"{sub}.{target_domain}"
                executor.submit(resolve_domain, full_domain, timeout)
    except Exception as e:
        print(f"{Colors.RED}[-] 子域名模块异常: {str(e)}{Colors.RESET}")


# -------------------------- 6. 端口扫描 --------------------------
def scan_port(host, port, timeout=2):
    """扫描单个端口是否开放"""
    sock = None
    try:
        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        sock.settimeout(timeout)
        result = sock.connect_ex((host, port))
        if result == 0:
            # 尝试获取服务信息
            service = "unknown"
            try:
                service = socket.getservbyport(port)
            except:
                pass
            print(f"{Colors.GREEN}[+] 端口开放  {host}:{port}  -> 服务: {service}{Colors.RESET}")
    except socket.timeout:
        print(f"{Colors.YELLOW}[!] 端口超时  {host}:{port}{Colors.RESET}")
    except Exception as e:
        print(f"{Colors.RED}[-] 端口扫描  {host}:{port}  异常 -> {str(e)}{Colors.RESET}")
    finally:
        if sock:
            sock.close()


def port_run(host, start_port, end_port, timeout=2):
    """批量端口扫描"""
    try:
        print(f"{Colors.YELLOW}[*] 端口扫描: {host}（{start_port}-{end_port}）{Colors.RESET}")

        # 验证端口范围
        if start_port < 1 or end_port > 65535 or start_port > end_port:
            print(f"{Colors.RED}[-] 无效的端口范围！请输入1-65535之间的有效范围{Colors.RESET}")
            return

        with ThreadPoolExecutor(max_workers=100) as executor:
            for port in range(start_port, end_port + 1):
                executor.submit(scan_port, host, port, timeout)

    except Exception as e:
        print(f"{Colors.RED}[-] 端口扫描模块异常: {str(e)}{Colors.RESET}")

# -------------------------- 命令行参数解析 --------------------------
def main():
    parser = argparse.ArgumentParser(description="智慧AI漏扫工具 v1.3")
    subparsers = parser.add_subparsers(help="选择功能模块", dest="sub_command")
    subparsers.required = True  # 强制选择模块

    # 1. MySQL参数
    mysql_parser = subparsers.add_parser("mysql", help="MySQL弱口令扫描")
    mysql_parser.add_argument("--host", "-H", default="127.0.0.1", help="目标主机（默认127.0.0.1）")
    mysql_parser.add_argument("--port", "-P", type=int, default=3306, help="目标端口（默认3306）")
    mysql_parser.add_argument("--user", "-u", type=argparse.FileType("r", encoding="utf-8"),
                              default="./user.dict", help="用户名字典（默认./user.dict）")
    mysql_parser.add_argument("--passwd", "-p", type=argparse.FileType("r", encoding="utf-8"),
                              default="./passwd.dict", help="密码字典（默认./passwd.dict）")

    # 2. Redis参数
    redis_parser = subparsers.add_parser("redis", help="Redis弱口令/未授权扫描")
    redis_parser.add_argument("--host", "-H", default="127.0.0.1", help="目标主机（默认127.0.0.1）")
    redis_parser.add_argument("--port", "-P", type=int, default=6379, help="目标端口（默认6379）")
    redis_parser.add_argument("--passwd", "-p", type=argparse.FileType("r", encoding="utf-8"),
                              default="./passwd.dict", help="密码字典（默认./passwd.dict）")

    # 3. IP参数
    ip_parser = subparsers.add_parser("ip", help="IP网段存活扫描")
    ip_parser.add_argument("--host", "-H", required=True, help="IP网段（如192.168.1.0/24，必填）")
    ip_parser.add_argument("--iface", "-f", default="以太网", help="ARP扫描网卡（默认以太网）")
    ip_parser.add_argument("--port", "-P", type=int, default=80, help="TCP检测端口（默认80）")

    # 4. SSH参数
    ssh_parser = subparsers.add_parser("ssh", help="SSH弱口令扫描")
    ssh_parser.add_argument("--host", "-H", default="127.0.0.1", help="目标主机（默认127.0.0.1）")
    ssh_parser.add_argument("--port", "-P", type=int, default=22, help="目标端口（默认22）")
    ssh_parser.add_argument("--user", "-u", type=argparse.FileType("r", encoding="utf-8"),
                            default="./user.dict", help="用户名字典（默认./user.dict）")
    ssh_parser.add_argument("--passwd", "-p", type=argparse.FileType("r", encoding="utf-8"),
                            default="./passwd.dict", help="密码字典（默认./passwd.dict）")

    # 5. 子域名参数（修复：完整编码格式+括号闭合）
    subdomain_parser = subparsers.add_parser("subdomain", help="子域名扫描")
    subdomain_parser.add_argument("--target", "-t", required=True, help="主域名（如baidu.com，必填）")
    subdomain_parser.add_argument("--dict", "-d", type=argparse.FileType("r", encoding="utf-8"),
                                  default="./subdomain.dict", help="子域名字典（默认./subdomain.dict）")
    subdomain_parser.add_argument("--timeout", "-to", type=int, default=3, help="解析超时时间（默认3秒）")

    # 6. 端口扫描参数
    port_parser = subparsers.add_parser("port", help="端口范围扫描")
    port_parser.add_argument("--host", "-H", required=True, help="目标主机（如192.168.1.1，必填）")
    port_parser.add_argument("--start", "-s", type=int, default=1, help="起始端口（默认1）")
    port_parser.add_argument("--end", "-e", type=int, default=1000, help="结束端口（默认1000）")
    port_parser.add_argument("--timeout", "-to", type=int, default=2, help="超时时间（默认2秒）")
    # 解析参数并执行
    args = parser.parse_args()

    if args.sub_command == "mysql":
        mysql_run(args.host, args.port, args.user.readlines(), args.passwd.readlines())
        args.user.close()
        args.passwd.close()

    elif args.sub_command == "redis":
        redis_run(args.host, args.port, args.passwd.readlines())
        args.passwd.close()

    elif args.sub_command == "ip":
        run_ip(args.host, args.iface, args.port)

    elif args.sub_command == "ssh":
        ssh_run(args.host, args.port, args.user.readlines(), args.passwd.readlines())
        args.user.close()
        args.passwd.close()

    elif args.sub_command == "subdomain":
        subdomain_run(args.target, args.dict.readlines(), args.timeout)
        args.dict.close()
    elif args.sub_command == "port":
        port_run(args.host, args.start, args.end, args.timeout)

if __name__ == "__main__":
    main()